import itertools

import numpy

from LFP import signal_utils
from projects.electrophysiology.view_data import frame_utils

def get_specgram_slice(signal, sampling_freq, time_resolution,
                       frequency_resolution, bath_signals, 
                       high_frequency_cutoff, target_frequency):
    '''
    Helper function for LFP_weighted_avg that calculates and returns
        the power vs time at a specific freqency of a signal.
    '''
    # find the part of the signal that could have been impacted by the
    #   stimulus being shown.
    Pxx, freqs, bins = signal_utils.specgram(signal, sampling_freq, 
                           time_resolution, frequency_resolution,
                           bath_signals=bath_signals,
                           high_frequency_cutoff=high_frequency_cutoff)

    # find the power vs time at the target frequency
    p_target_full, bins = signal_utils.specgram_slice(Pxx, freqs, bins,
                                                      target_frequency)
    return p_target_full, bins

def LFP_weighted_avg(signals, sampling_freq, pulse_index, frames,
                     frame_times_list, target_frequency, tau, 
                     noise_signals=[], 
                     time_resolution=1.0/120.0, # half a frame time 
                     frequency_resolution=1.0):
    """
    This function computes the LFP weighted average and returns back a 2-D 
       array.    
    INPUTS:
        signals           : list of voltage traces
        sampling_freq     : sampling frequency at which the data were 
                             taken (in Hz)
        pulse_index       : index of corresponding to when the stimulus was
                             triggered to begin.
        frames            : visual stimulus frames
        frame_times_list  : for each signal, a list of times when frames were
                             displayed. (in seconds)
        target_frequency  : frequency in the LFP used in constructing 
                             the Weighted average (in Hz)
        tau               : amount of time between frame display and when the
                             lfp is affected (in seconds)
        --kwargs--
        noise_signals=[]            : list of signals aquired without stimulus 
                                       used to eliminate noise in the LFP.
        time resolution=1.0/120     : the minimum time resolution of the 
                                       calculated spectrograms (in seconds)
        frequency resolution=1.0    : the minimum frequency resolution of the
                                       calculated psd and spectrograms (in Hz)
    Returns:
        image           : an image like those in <frames> representing the 
                           LFP weighted average stimulus.
        power_norm      : power normalization factor
    """
    # don't bother calculating the spectrogram above 1.5 x target_freq.
    high_frequency_cutoff = 3*target_frequency 

    results = []
    for signal, frame_times in zip(signals, frame_times_list):
        # find the power vs. time for a given target freqency
        p_target_full, bins = get_specgram_slice(signal, sampling_freq, 
                                                 time_resolution, 
                                                 frequency_resolution,
                                                 noise_signals, 
                                                 high_frequency_cutoff,
                                                 target_frequency)
        bins = numpy.array(bins)

        # With the power vs. time of the entire signal, slice down to
        #   just the part of the power when the stimulus could have 
        #   affected it.
        begin_time = pulse_index/sampling_freq + tau
        end_time = begin_time + frame_times[-1]
        p_target_high = p_target_full[bins>begin_time]
        bins_high = bins[bins>begin_time]
        # p_target = power when frames were able to affect it.
        p_target = p_target_high[bins_high<=end_time]
        bins_target = bins_high[bins_high<=end_time]
        # p_bins = times relative to stimulus onset.
        p_bins = bins_target - begin_time 

        # normalize the power so that it sums to 1.0
        power_norm = numpy.sum(p_target)
        norm_p = p_target/power_norm

        LFP_weighted_stimuli = numpy.zeros(frames[0].shape, 
                                           dtype=numpy.float64)
        for time, power in itertools.izip(p_bins, norm_p):
            LFP_weighted_stimuli += (frame_utils.get_stim(time, frames, 
                                                          frame_times,
                                                          None, None) * power)
        results.append(LFP_weighted_stimuli)
    # Average out the TRIAL aspect
    return numpy.average(numpy.array(results), axis=0), power_norm
    
